package ai.mspbots.gateway.filter;

import co.elastic.apm.api.Transaction;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.util.UriComponentsBuilder;
import reactor.core.publisher.Mono;
import co.elastic.apm.api.ElasticApm;

import java.net.URI;

@Component
public class MdcFilter implements GlobalFilter {

    private static final Logger logger = LoggerFactory.getLogger(MdcFilter.class);

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        Transaction transaction = ElasticApm.currentTransaction();
        try {
            // 设置 MDC 值
            String traceId = transaction.getTraceId();
            String transactionId = transaction.getId();
            MDC.put("trace.id", traceId);
            MDC.put("transaction.id", transactionId);
            logger.info("config trace.id: {}", traceId);
            logger.info("config transaction.id: {}", transactionId);

            URI originUri = exchange.getRequiredAttribute(ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR);
            URI uri = UriComponentsBuilder.fromUri(originUri)
                    .build().toUri();

            ServerHttpRequest request = exchange.getRequest();
            ServerHttpRequest modifiedRequest = request.mutate()
                    .headers(httpHeaders -> {
                        httpHeaders.add("trace.id", traceId);
                        httpHeaders.add("transaction.id", transactionId);
                    })
                    .uri(uri)
                    .build();
            exchange.getAttributes().put(ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR, uri);
            ServerWebExchange webExchange = exchange.mutate().request(modifiedRequest).build();
            return chain.filter(webExchange).doFinally(signalType -> {
                // 清理 MDC
                MDC.clear();
                // 结束事务
                transaction.end();
            });
        } catch (Exception e) {
            // 在发生异常时记录日志或处理异常
            MDC.clear();
            transaction.captureException(e);
            transaction.end();
            throw e; // 重新抛出异常
        }
    }
}